{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Experiment with MiniWoB++ environment"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# if True, run MiniWoB++ in headless mode\n",
    "headless = False\n",
    "\n",
    "# number of episodes to run per task\n",
    "num_episodes = 1\n",
    "\n",
    "# set to True to use the expert samples\n",
    "ALLOW_DEMO = True\n",
    "# set to False to directly use generated plan without syntax check\n",
    "CODE_CHECK = False\n",
    "# 'gpt-3.5' / 'gpt-3' / 'davinci-003'\n",
    "GPT_MODEL = 'davinci-003'\n",
    "\n",
    "# max number of trials for each task\n",
    "NUM_TRY = 20\n",
    "# after NUM_TRY_RESET trials, the agent will try to start from step 1.\n",
    "NUM_TRY_RESET = 10\n",
    "\n",
    "# list of tasks where in-loop refinement is needed.\n",
    "tasks_with_ask = [\n",
    "    'simple-algebra',\n",
    "    'click-shades',\n",
    "    'click-shape',\n",
    "    'identify-shape',\n",
    "    'count-shape',\n",
    "    'click-checkboxes-soft',\n",
    "    'search-engine',\n",
    "    'tic-tac-toe',\n",
    "    'guess-number',\n",
    "    'terminal',\n",
    "]\n",
    "\n",
    "# list of tasks where feedback is available.\n",
    "tasks_with_feedback = [\n",
    "    \"search-engine\",\n",
    "    \"tic-tac-toe\",\n",
    "    \"email-inbox-forward-nl-turk\",\n",
    "    \"terminal\",\n",
    "    \"login-user-popup\",\n",
    "    \"guess-number\",\n",
    "    \"email-inbox-forward-nl\",\n",
    "    \"email-inbox\",\n",
    "    \"email-inbox-nl-turk\"\n",
    "]\n",
    "\n",
    "# save to files\n",
    "LOG_FILE = 'interaction_log.txt'\n",
    "PR_FILE = 'Prompt_Response.txt'"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import argparse\n",
    "import random\n",
    "\n",
    "import numpy as np\n",
    "import os\n",
    "import openai\n",
    "import computergym\n",
    "import gym\n",
    "\n",
    "from selenium import webdriver\n",
    "from selenium.webdriver.common.by import By\n",
    "from selenium.webdriver.common.action_chains import ActionChains\n",
    "from datetime import datetime\n",
    "import logging\n",
    "\n",
    "logging.basicConfig(level=logging.WARNING)\n",
    "\n",
    "import urllib3\n",
    "urllib3.disable_warnings()\n",
    "\n",
    "from computergym.miniwob.miniwob_interface.action import (\n",
    "    MiniWoBType,\n",
    "    MiniWoBElementClickId,\n",
    "    MiniWoBElementClickXpath,\n",
    "    MiniWoBElementClickOption,\n",
    "    MiniWoBMoveXpath,\n",
    ")\n",
    "import re\n",
    "from selenium.webdriver.common.keys import Keys\n",
    "import json\n",
    "import time\n",
    "import openai\n",
    "from pathlib import Path\n",
    "\n",
    "import io\n",
    "import traceback\n",
    "\n",
    "import sys\n",
    "sys.stdout.flush()\n",
    "\n",
    "openai.api_key = os.environ[\"OPENAI_API_KEY\"]\n",
    "\n",
    "def save_to_file(directory, filename, content):\n",
    "    if not os.path.exists(directory):\n",
    "        os.makedirs(directory)\n",
    "        \n",
    "    file_path = os.path.join(directory, filename)\n",
    "    with open(filename, 'a') as file:\n",
    "        file.write(content + '\\n')\n",
    "\n",
    "def clear_file_content(filename):\n",
    "    with open(filename, 'w') as file:\n",
    "        pass\n",
    "\n",
    "def dict_to_traj(d):\n",
    "    traj = ''\n",
    "    for i, a in enumerate(d['actions']):\n",
    "        traj += f'Act {i}: {a}\\nObs {i}: {d[\"observations\"][i]}\\n'\n",
    "    return traj\n",
    "\n",
    "def extract_answers(text, markers):\n",
    "    answers = []\n",
    "    for i, marker in enumerate(markers):\n",
    "        if i < len(markers) - 1:\n",
    "            next_marker = markers[i + 1]\n",
    "            pattern = fr\"{re.escape(marker)}\\s*([\\s\\S]*?)\\n*{re.escape(next_marker)}\"\n",
    "        else:\n",
    "            pattern = fr\"{re.escape(marker)}\\s*([\\s\\S]*)\" \n",
    "        answer = re.search(pattern, text)\n",
    "        if answer:\n",
    "            answer = answer.group(1).strip()\n",
    "            answers.append(answer)\n",
    "        else:\n",
    "            answers.append('Not found')\n",
    "    return answers\n",
    "\n",
    "def eliminate_content_after(text, eliminate_from):\n",
    "    lines = text.split('\\n')\n",
    "    result = []\n",
    "    \n",
    "    for line in lines:\n",
    "        if eliminate_from in line:\n",
    "            break\n",
    "        result.append(line)\n",
    "    \n",
    "    return '\\n'.join(result)\n",
    "\n",
    "def get_error_step(error_message):\n",
    "    pattern = r'\\[Step (\\d+)\\]'\n",
    "    match = re.search(pattern, error_message)\n",
    "    if match:\n",
    "        return int(match.group(1))\n",
    "    else:\n",
    "        return None\n",
    "    \n",
    "def get_first_digit(input_string):\n",
    "    pattern = r'\\d'\n",
    "    match = re.search(pattern, input_string)\n",
    "    if match:\n",
    "        return int(match.group())\n",
    "    else:\n",
    "        return None\n",
    "\n",
    "datetime = datetime.now().strftime(\"%m%d-%H%M%S\")\n",
    "save_to_folder = f'./r_{GPT_MODEL}_{datetime}/'\n",
    "LOG_FILE =  save_to_folder + LOG_FILE\n",
    "PR_FILE =  save_to_folder + PR_FILE\n",
    "\n",
    "save_to_file(save_to_folder, PR_FILE, 'Start: ' + str(datetime) + '\\n')\n",
    "save_to_file(save_to_folder, LOG_FILE, 'Start: ' + str(datetime) + '\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def ask(prompt):\n",
    "    prompt_chat = [\n",
    "            {\"role\": \"user\", \"content\": prompt.strip()},\n",
    "        ]\n",
    "    cnt = 0\n",
    "    while True:\n",
    "        try:\n",
    "            if GPT_MODEL == 'gpt-3.5':\n",
    "                response = openai.ChatCompletion.create(\n",
    "                    model=\"gpt-3.5-turbo\",\n",
    "                    messages=prompt_chat,\n",
    "                    temperature=0,\n",
    "                    max_tokens=1400,\n",
    "                    top_p=1,\n",
    "                    frequency_penalty=0.0,\n",
    "                    presence_penalty=0.0,\n",
    "                    )\n",
    "                answer = response[\"choices\"][0][\"message\"]['content'].strip()\n",
    "                save_to_file(save_to_folder,\n",
    "                             PR_FILE, \n",
    "                            'Prompt: \\n' +\n",
    "                            prompt +\n",
    "                            '\\nResponse: \\n' +\n",
    "                            answer + '\\n' + '='*20 + '\\n')\n",
    "        \n",
    "                return answer\n",
    "            elif GPT_MODEL == 'gpt-3' or GPT_MODEL == 'davinci-003':\n",
    "                response = openai.Completion.create(\n",
    "                    model=\"text-davinci-002\" if GPT_MODEL == 'gpt-3' else \"text-davinci-003\",\n",
    "                    prompt=prompt,\n",
    "                    temperature=0,\n",
    "                    max_tokens=1000,\n",
    "                    top_p=1,\n",
    "                    frequency_penalty=0.0,\n",
    "                    presence_penalty=0.0,\n",
    "                    )\n",
    "                answer = response[\"choices\"][0][\"text\"].strip()\n",
    "                save_to_file(save_to_folder,\n",
    "                            PR_FILE,\n",
    "                            'Prompt: \\n' +\n",
    "                            prompt +\n",
    "                            '\\nResponse: \\n' +\n",
    "                            answer + '\\n' + '='*20 + '\\n')\n",
    "                return answer\n",
    "            else:\n",
    "                raise Exception('Wrong GPT_MODEL')\n",
    "        except openai.error.RateLimitError as e:\n",
    "            retry_after = 3\n",
    "            print(f\"Rate limit error: {e}. Retrying in {retry_after} seconds.\")\n",
    "            time.sleep(retry_after)\n",
    "        except openai.error.InvalidRequestError as e:\n",
    "            eliminate_context = \"# for example: You have a list of receptacles, and you want to sort them by the likelihood of a soapbar appearing in them. You can do this by asking the assistant:\\nreceptacles = ['countertop 1', 'garbagecan 1', 'sinkbasin 2', 'sinkbasin 1', 'toilet 1', 'toiletpaperhanger 1', 'towelholder 1']\\nanswer = ask(f'Sort the list of receptacles, starting from the one a soapbar is most likely to appear: {receptacles}. You should return a Python list.')\\n# answer = ['sinkbasin 1', 'sinkbasin 2', 'countertop 1', 'towelholder 1', 'toiletpaperhanger 1', 'garbagecan 1', 'toilet 1']\"\n",
    "            prompt = prompt.replace(eliminate_context, '')\n",
    "            prompt_chat = [\n",
    "                {\"role\": \"user\", \"content\": prompt.strip()},\n",
    "            ]\n",
    "            print(f\"Exceed max: {e}.\")\n",
    "            cnt += 1\n",
    "            if cnt > 3:\n",
    "                return 'Exceed max limit. Tried 3 times. Skip this one.'\n",
    "        except openai.error.APIError:\n",
    "            cnt += 1\n",
    "            if cnt > 3:\n",
    "                return 'APIError. Tried 3 times. Skip this one.'\n",
    "        except Exception as e:\n",
    "            print(f\"An unexpected error occurred: {e}\")\n",
    "            cnt += 1\n",
    "            if cnt > 3:\n",
    "                return 'Error. Tried 3 times. Skip this one.'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def capture_output(func, agent, step=1):\n",
    "    # Store the original standard output and standard error\n",
    "    original_stdout = sys.stdout\n",
    "    original_stderr = sys.stderr\n",
    "\n",
    "    # Redirect the standard output and error to in-memory file-like objects\n",
    "    temp_stdout = io.StringIO()\n",
    "    temp_stderr = io.StringIO()\n",
    "    sys.stdout = temp_stdout\n",
    "    sys.stderr = temp_stderr\n",
    "\n",
    "    checkpoint = None\n",
    "    # Run the function and capture exceptions\n",
    "    try:\n",
    "        func(agent, start_from=step)\n",
    "    except Exception as e:\n",
    "        traceback.print_exc()\n",
    "        checkpoint = sys.exc_info()[2].tb_next.tb_frame.f_locals\n",
    "\n",
    "    # Restore the original standard output and error\n",
    "    sys.stdout = original_stdout\n",
    "    sys.stderr = original_stderr\n",
    "\n",
    "    # Get the output and error messages as strings\n",
    "    output_string = temp_stdout.getvalue()\n",
    "    error_string = temp_stderr.getvalue()\n",
    "\n",
    "    print(output_string)\n",
    "    print(error_string)\n",
    "    return error_string, checkpoint, output_string + error_string\n",
    "\n",
    "def modify_header(checkpoint):\n",
    "    if not checkpoint:\n",
    "        return 'def solution(agent, start_from=1):'\n",
    "    load_checkpoint = ''\n",
    "    skip_vars = ['agent', 'start_from']\n",
    "    for k,v in checkpoint.items():\n",
    "        if k not in skip_vars:\n",
    "            if type(v) == str:\n",
    "                load_checkpoint += f', {k}=\"{v}\"'\n",
    "            else:\n",
    "                load_checkpoint += f', {k}={v}'\n",
    "    header = f'def solution(agent, start_from{load_checkpoint}):'\n",
    "    return header"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Define the agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# define the Agent\n",
    "\n",
    "# adapted from rci\n",
    "def get_html_state(env, states):\n",
    "    if not states[0]:\n",
    "        return None\n",
    "    # return states[0]._dom_elements\n",
    "    extra_html_task = [\n",
    "        \"click-dialog\",\n",
    "        \"click-dialog-2\",\n",
    "        \"use-autocomplete\",\n",
    "        \"choose-date\",\n",
    "    ]\n",
    "    html_body = states[0].html_body\n",
    "    if env in extra_html_task:\n",
    "        html_body += states[0].html_extra\n",
    "    # return bs(html_body, 'html.parser')\n",
    "    return html_body\n",
    "\n",
    "class Agent:\n",
    "    def __init__(self, env_name):\n",
    "        self.env_name = env_name\n",
    "        self.env = gym.make(\"MiniWoBEnv-v0\", env_name=env_name, headless=headless)\n",
    "        observation = self.env.reset(seeds=[random.random()], record_screenshots=False)\n",
    "        self.description = observation[0].utterance\n",
    "        self.current_state = get_html_state(env_name, observation)\n",
    "        self.interaction_history = {'actions': [], 'observations': []}\n",
    "        self.is_success = False\n",
    "\n",
    "    def reset(self):\n",
    "        observation = self.env.reset(seeds=[random.random()], record_screenshots=False)\n",
    "        self.description = observation[0].utterance\n",
    "        self.current_state = get_html_state(self.env_name, observation)\n",
    "        self.interaction_history = {'actions': [], 'observations': []}\n",
    "        self.is_success = False\n",
    "        self.add_to_history('Initial state', str(self.current_state))\n",
    "\n",
    "    def close(self):\n",
    "        self.env.close()\n",
    "\n",
    "    def __del__(self): \n",
    "        print(f\"Current task {self.env_name} is done.\")\n",
    "\n",
    "    def interact_with_env(self, action):\n",
    "        observation, reward, done, _ = self.env.step([action])\n",
    "        # print(observation, reward, done)\n",
    "        html_obs, reward, done = get_html_state(self.env_name, observation), reward[0], all(done)\n",
    "        # html_obs = observation[0]._dom_elements\n",
    "        return html_obs, reward, done\n",
    "    \n",
    "    # Note down the history of interactions with the environment\n",
    "    def add_to_history(self, action, observation):\n",
    "        self.interaction_history['actions'].append(action)\n",
    "        self.interaction_history['observations'].append(observation)\n",
    "        \n",
    "    # Get an observation from the environment after performing an action, and add it to the history\n",
    "    def observation(self, action):\n",
    "        # add some sleep time between each step.\n",
    "        time.sleep(0.3)\n",
    "        html_obs, reward, done = self.interact_with_env(action)\n",
    "        self.current_state = html_obs\n",
    "        self.add_to_history(action, html_obs)\n",
    "        print(f'Act: {action}, reward: {reward}')\n",
    "        if done or reward != 0:\n",
    "            if reward > 0:\n",
    "                self.is_success = True\n",
    "            print('Done. Success:', reward)\n",
    "            return 'DONE'\n",
    "        return html_obs\n",
    "    \n",
    "    # Here are the admissible actions:\n",
    "    # Action: type a string into the input box\n",
    "    # this function returns the html after the action\n",
    "    def type(self, characters: str):\n",
    "        action = MiniWoBType(characters)\n",
    "        return self.observation(action)\n",
    "\n",
    "    # Actions: press a key on the keyboard, the input can be one of the following:\n",
    "    # enter, space, arrow_left, arrow_right, arrow_up, arrow_down, backspace\n",
    "    # this function returns the html after the action\n",
    "    def press_key(self, key: str):\n",
    "        if key == 'enter':\n",
    "            miniwob_key = '\\n'\n",
    "        elif key == 'space':\n",
    "            miniwob_key = ' '\n",
    "        elif key == 'arrow_left':\n",
    "            miniwob_key = Keys.LEFT\n",
    "        elif key == 'arrow_right':\n",
    "            miniwob_key = Keys.RIGHT\n",
    "        elif key == 'arrow_up':\n",
    "            miniwob_key = Keys.UP\n",
    "        elif key == 'arrow_down':\n",
    "            miniwob_key = Keys.DOWN\n",
    "        elif key == 'backspace':\n",
    "            miniwob_key = Keys.BACKSPACE\n",
    "        action = MiniWoBType(miniwob_key)\n",
    "        return self.observation(action)\n",
    "        \n",
    "    # Action: click an option HTML element in a list with an XPath\n",
    "    # this function returns the html after the action\n",
    "    def click_option(self, xpath: str):\n",
    "        action = MiniWoBElementClickOption(xpath)\n",
    "        return self.observation(action)\n",
    "\n",
    "    # Action: click an HTML element with an XPath\n",
    "    # this function returns the html after the action\n",
    "    def click_xpath(self, xpath: str):\n",
    "        action = MiniWoBElementClickXpath(xpath)\n",
    "        return self.observation(action)\n",
    "\n",
    "    # Action: click an HTML element with an ID\n",
    "    # this function returns the html after the action\n",
    "    def click_id(self, element_id: str):\n",
    "        action = MiniWoBElementClickId(element_id)\n",
    "        return self.observation(action)\n",
    "\n",
    "    # Action: move the mouse cursor on an HTML element with an XPath\n",
    "    # this function returns the html after the action\n",
    "    def move_mouse_on(self, xpath: str):\n",
    "        action = MiniWoBMoveXpath(xpath)\n",
    "        return self.observation(action)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Closed-loop simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from prompt import get_solution_prompt_CL_with_ask, get_solution_prompt_CL_no_ask, feedback_fix_prompt_with_ask, feedback_fix_prompt_no_ask, code_check_prompt, example_list, get_start_from_prompt\n",
    "import textwrap\n",
    "\n",
    "results = {k: [] for k in tasks_with_feedback}\n",
    "skills = {k: [] for k in tasks_with_feedback}\n",
    "\n",
    "for idx, task in enumerate(tasks_with_feedback):\n",
    "    env_name = task\n",
    "\n",
    "    if env_name in tasks_with_ask:\n",
    "        get_solution_prompt = get_solution_prompt_CL_with_ask\n",
    "        feedback_fix_prompt = feedback_fix_prompt_no_ask\n",
    "    else:\n",
    "        get_solution_prompt = get_solution_prompt_CL_no_ask\n",
    "        feedback_fix_prompt = feedback_fix_prompt_with_ask\n",
    "\n",
    "    print(f'Task: {env_name}')\n",
    "    agent = Agent(env_name)\n",
    "    for episode in range(num_episodes):\n",
    "        agent.reset()\n",
    "        terminal_output = ''\n",
    "        print(f'Task: {agent.description}')\n",
    "        initial_state = str(agent.current_state)\n",
    "\n",
    "        example = ''\n",
    "        if ALLOW_DEMO:\n",
    "            example = example_list[env_name]\n",
    "\n",
    "        prompt = get_solution_prompt\\\n",
    "                .replace('<initial_state>', initial_state)\\\n",
    "                .replace('<task>', agent.description)\\\n",
    "                .replace('<example>', example)\n",
    "        response = ask(prompt)\n",
    "        \n",
    "        \n",
    "        # refine internally                \n",
    "        solution_func = '''\n",
    "        def solution(agent, start_from=1):\n",
    "        <solution>\n",
    "        '''.strip().replace('<solution>', response) if not response.startswith('def solution(agent, start_from=1):') else response\n",
    "\n",
    "        solution_func = eliminate_content_after(solution_func, '# Now complete the function')\n",
    "        solution_func = eliminate_content_after(solution_func, 'solution(agent)')\n",
    "\n",
    "\n",
    "        if CODE_CHECK:\n",
    "            prompt = code_check_prompt\\\n",
    "                    .replace('<solution_func>', solution_func)\n",
    "            response = ask(prompt)\n",
    "            answers = extract_answers(response, ['[Decision]', '[Revised code]'])\n",
    "\n",
    "            if 'def solution' in answers[1]:\n",
    "                solution_func = answers[1].strip('```')\\\n",
    "                    .replace('[Revised code]', '')\\\n",
    "                    .replace('without any other words.', '')\\\n",
    "                    .replace('Revised code:', '').strip()\\\n",
    "                    .replace('```', '')\n",
    "\n",
    "        solution_func = solution_func.replace('agent.ask', 'ask').replace('return', '# return').replace('solution(agent)','')\n",
    "\n",
    "        # Replace 'solution(agent)' with '' but leave 'def solution(agent)' untouched\n",
    "        pattern = r'(?<!def )solution\\(agent\\)'\n",
    "        solution_func = re.sub(pattern, '', solution_func)\n",
    "        start_num = 1\n",
    "        for num_try in range(NUM_TRY):\n",
    "            if num_try < NUM_TRY_RESET:\n",
    "                step = start_num if start_num else 1\n",
    "            else:\n",
    "                step = 1\n",
    "            print('start_from_step:', start_num)\n",
    "            # execute the solution function\n",
    "            def_error = False\n",
    "            try:\n",
    "                print(solution_func)\n",
    "                exec(solution_func)\n",
    "            except Exception as e:\n",
    "                error_msg = str(e)\n",
    "                error_string = str(e)\n",
    "                checkpoint = None\n",
    "                def_error = True\n",
    "\n",
    "            if not def_error:\n",
    "                error_string, checkpoint, output_string = capture_output(solution, agent, step)\n",
    "                terminal_output += output_string\n",
    "                if error_string:\n",
    "                    error_msg = error_string.split('\\n')[4:]\n",
    "                    error_msg = '\\n'.join(error_msg)\n",
    "                else:\n",
    "                    error_msg = 'You executed the solution function successfully but the task is not completed. Please check your solution function.'\n",
    "            \n",
    "            start_num = None\n",
    "\n",
    "            if agent.is_success or 'DONE' in terminal_output or 'Done. Success:' in terminal_output:\n",
    "                break\n",
    "\n",
    "            prev_solution_func = solution_func\n",
    "\n",
    "            if 'NoneType' in error_string:\n",
    "                break\n",
    "\n",
    "            # refine based on environment feedback\n",
    "            prompt = feedback_fix_prompt\\\n",
    "                    .replace('<solution_func>', prev_solution_func)\\\n",
    "                    .replace('<task>', agent.description)\\\n",
    "                    .replace('<feedback>', error_msg)\n",
    "            response = ask(prompt)\n",
    "            solution_func = '''\n",
    "            def solution(agent, start_from=1):\n",
    "            <solution>\n",
    "            '''.strip().replace('<solution>', textwrap.indent(response, \"\\t\")) if not response.startswith('def solution(agent, start_from=1):') else response\n",
    "            solution_func = solution_func.replace('agent.ask', 'ask').replace('return', '# return').replace('solution(agent)','')\n",
    "\n",
    "            prompt = get_start_from_prompt\\\n",
    "                    .replace('<previous_solution>', prev_solution_func)\\\n",
    "                    .replace('<revised_solution>', solution_func)\n",
    "            response = ask(prompt)\n",
    "            start_num = get_first_digit(response)\n",
    "            # temporarily not needed\n",
    "            # solution_func = solution_func.replace('def solution(agent, start_from=1):', modify_header(checkpoint))\n",
    "\n",
    "        results[task].append(agent.is_success)\n",
    "        \n",
    "\n",
    "        save_to_file(save_to_folder, LOG_FILE, \n",
    "                    f'Task {episode+1}: {agent.env_name}\\n' + \\\n",
    "                    agent.description + '\\n' + \\\n",
    "                    dict_to_traj(agent.interaction_history) + '\\n' + \\\n",
    "                    f'Success: {agent.is_success}\\n')\n",
    "        save_to_file(save_to_folder, LOG_FILE, f'results {results}')\n",
    "        save_to_file(save_to_folder, LOG_FILE, f'------------\\n')\n",
    "\n",
    "        print(f'Task: {agent.env_name}, Episode {episode+1}, Success: {agent.is_success}')\n",
    "        print('-'*20)\n",
    "\n",
    "    agent.close()\n",
    "    del agent\n",
    "    print(f'Task {idx}: {env_name}, Success rate: {np.mean(results[env_name])}')\n",
    "    print('='*20)\n",
    "save_to_file(save_to_folder, LOG_FILE, f'Final results {results}')\n",
    "save_to_file(save_to_folder, LOG_FILE, f'Final skills {skills}')\n",
    "print(f'Final results {results}')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dl_env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
